import numpy as np
from pprint import pprint
from collections import deque

easy = [
    #        |        |
    [0, 0, 0, 0, 4, 0, 0, 0, 0],
    [0, 3, 0, 0, 1, 7, 9, 0, 0],
    [1, 9, 0, 3, 0, 0, 0, 4, 2],
    #--------------------------
    [0, 2, 0, 6, 3, 4, 0, 0, 8],
    [0, 5, 0, 2, 8, 1, 0, 7, 0],
    [3, 0, 0, 7, 5, 9, 0, 1, 0],
    #--------------------------
    [2, 1, 0, 0, 0, 3, 0, 8, 4],
    [0, 0, 4, 1, 7, 0, 0, 2, 0],
    [0, 0, 0, 0, 2, 0, 0, 0, 0],
    #        |        |
    ]

hardest = [
    #        |        |
    [0, 8, 0, 0, 0, 5, 0, 9, 0],
    [0, 0, 7, 4, 2, 0, 0, 0, 3],
    [0, 3, 0, 0, 0, 0, 0, 0, 1],
    #--------------------------
    [0, 0, 9, 0, 0, 0, 5, 0, 0],
    [0, 5, 6, 0, 4, 0, 1, 8, 0],
    [0, 0, 4, 0, 0, 0, 9, 0, 0],
    #--------------------------
    [7, 0, 0, 0, 0, 0, 0, 2, 0],
    [2, 0, 0, 0, 1, 8, 4, 0, 0],
    [0, 4, 0, 9, 0, 0, 0, 1, 0],
    #        |        |
    ]

n35 = [
    #        |        |
    [0, 5, 0, 0, 0, 9, 8, 0, 6],
    [0, 0, 0, 0, 0, 0, 0, 0, 0],
    [7, 0, 0, 8, 0, 0, 4, 9, 2],
    #--------------------------
    [4, 9, 0, 0, 0, 2, 0, 0, 3],
    [0, 3, 0, 0, 0, 0, 0, 2, 0],
    [8, 0, 0, 3, 0, 7, 0, 4, 5],
    #--------------------------
    [5, 8, 3, 0, 0, 1, 0, 0, 7],
    [0, 0, 0, 0, 0, 0, 0, 0, 0],
    [6, 0, 1, 7, 0, 0, 0, 3, 0],
    #        |        |
    ]

n10 = [ #unsolved
    #        |        |
    [0, 0, 0, 7, 0, 0, 0, 1, 0],
    [0, 0, 0, 0, 2, 0, 8, 0, 3],
    [0, 2, 0, 9, 0, 0, 0, 0, 5],
    #--------------------------
    [5, 0, 0, 3, 0, 0, 2, 0, 0],
    [1, 8, 0, 0, 7, 0, 0, 4, 6],
    [0, 0, 2, 0, 0, 6, 0, 0, 1],
    #--------------------------
    [8, 0, 0, 0, 0, 7, 0, 5, 0],
    [4, 0, 7, 0, 9, 0, 0, 0, 0],
    [0, 5, 0, 0, 0, 4, 0, 0, 0],
    #        |        |
    ]

def unique(n, a):
    return len(a[a==n]) == 1

class PySudoku(object):
    def __init__(self, grid, w=3, h=3):
        assert len(grid) == w*h and len(grid[0]) == w*h
        self.w, self.h = w, h
        self.size = w * h
        self.grid = np.array(grid, dtype=np.int8)
        self.tries = 0
        self.solved = self.__solve(self.grid)

    def __solve(self, grid):
        found = True
        while found and 0 in grid:
            while found and 0 in grid:
                found = False
                candidates = np.zeros((self.size, self.size, self.size), dtype=np.int8)
                for row, col in np.ndindex(self.size, self.size):
                    if not grid[row, col]:
                        for n in xrange(1, self.size + 1):
                            if (n not in grid[row, :] and
                                n not in grid[:, col] and
                                n not in grid[
                                    row // self.h * self.w : (row // self.h + 1) * self.w ,
                                    col // self.w * self.h : (col // self.w + 1) * self.h
                                ]): candidates[row, col, n-1] = n
                        cand = (candidates[row, col] > 0).sum()
                        if cand == 1:
                            grid[row, col] = candidates[row, col].max()
                            found = True
                        elif cand == 0: raise ValueError("Bad grid: no candidates in"
                                                " cell [%i, %i]" % (row, col))
            if 0 in grid:
                for row, col in np.ndindex(self.size, self.size):
                    if not grid[row, col]:
                        for n in candidates[row, col][candidates[row, col].nonzero()]:
                            if (unique(n, candidates[row, :]) or
                                unique(n, candidates[:, col]) or
                                unique(n, candidates[
                                    row // self.h * self.w : (row // self.h + 1) * self.w ,
                                    col // self.w * self.h : (col // self.w + 1) * self.h
                                ])):
                                    grid[row, col] = n
                                    found = True

        if 0 in grid: #dead point start guessing
            scores = []
            for row, col in np.ndindex(self.size, self.size):
                if not grid[row, col]:
                    n = (candidates[row, col] > 0).sum() + (grid[row, :] == 0).sum()\
                        + (grid[:, col] == 0).sum() + (grid[
                                row // self.h * self.w : (row // self.h + 1) * self.w ,
                                col // self.w * self.h : (col // self.w + 1) * self.h
                            ] == 0).sum()
                    scores.append((n, row, col))
            for n, row, col in sorted(scores):
                for cand in candidates[row, col][candidates[row, col] > 0]:
                    self.tries += 1
                    newGrid = grid.copy()
                    try:
                        newGrid[row, col] = cand
                        newGrid = self.__solve(newGrid)
                    except ValueError:
                        pass
                    else:
                        grid = newGrid
                        break
                if 0 not in grid: break
            else:
                raise ValueError("Don't know how to solve it after %i tries!" % self.tries)
            return grid
        else: return grid

    def __str__(self):
        box = self.size * self.h
        l = (self.w * 2 + 1)
        crosses = ("-" * l + "+") * (self.h - 1) + "-" * l
        dashes = ("-" * l * self.h)
        sep = "|\n|%s|\n| " % crosses
        row = deque(("\n+-%s-+\n| " % dashes,))
        for i, n in enumerate(self.solved.flat):
            if i % box == 0:
                if i: row.append(sep)
            elif i % self.size == 0: row.append("|\n| ")
            elif i % self.w == 0: row.append("| ")
            row.append(n and "%i " % n or '  ')
        row.append("|\n+-%s-+\n" % dashes)
        return ''.join(row)

    def __repr__(self):
        return "<PySudoku instance %ix%i at %s>" % (self.size, self.size, hex(id(self)))

if __name__ == "__main__":
    s = PySudoku(hardest)
    print s
    print s.tries
